package com.coder.algorithm.offer_2;

import java.util.Arrays;
import java.util.Locale;

/**
 * 最小的 k 个数
 * <p>
 * 解法二：借助最大堆（或者红黑树）
 *
 * @author yuhushuan
 * @date 2020/6/25 00:20
 */
public class _40_1_KLeastNumbers {

    /**
     * 时间复杂度 O(n)
     *
     * @param array
     * @param k
     */
    private static int[] getLeastNumbers(int[] array, int k) {
        if (array == null || array.length == 0 || k <= 0 || k > array.length) {
            return null;
        }
        int start = 0;
        int end = array.length - 1;
        int index = partition(array, start, end);
        while (index != k - 1) {// 可能存在 k == array.length 的情况
            if (index > k - 1) {
                end = index - 1;
                index = partition(array, start, end);
            } else {
                start = index + 1;
                index = partition(array, start, end);
            }
        }
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = array[i];
        }
        return result;
    }

    private static int partition(int[] array, int start, int end) {
        if (array == null || start < 0 || end < start) {
            throw new IllegalArgumentException("Invalid Argument");
        }
        int index = (int) (start + Math.random() * (end - start + 1));
        swap(array, index, end);// 基准位于末尾

        int small = start - 1;
        for (int i = start; i <= end; i++) {
            if (array[i] <= array[end]) {
                small++;
                // 此时 array[small] 大于 temp
                if (i > small) {
                    swap(array, i, small);
                }
            }
        }
        return small;
    }

    private static void swap(int[] array, int i, int j) {
        int temp = array[i];
        array[i] = array[j];
        array[j] = temp;
    }

    public static void main(String[] args) {
        // 功能测试
        // 数组中没有相同数字
        print(new int[]{4, 5, 1, 6, 2, 7, 3, 8}, 4);
        // 数组中有相同数字
        print(new int[]{2, 2, 2, 7, 8, 1, 2, 2}, 4);

        // 边界值测试
        // k = 1
        print(new int[]{4, 5, 1, 6, 2, 7, 3, 8}, 1);
        // k = 数组长度
        print(new int[]{4, 5, 1, 6, 2, 7, 3, 8}, 8);

        // 特殊输入测试
        print(new int[]{4, 5, 1, 6, 2, 7, 3, 8}, 0);
        print(new int[]{4, 5, 1, 6, 2, 7, 3, 8}, 10);
        print(null, 2);
    }

    private static void print(int[] array, int k) {
        System.out.println(String.format(Locale.getDefault(),
                "数组：%s\n最小的 %d 个数为：\n%s\n",
                Arrays.toString(array), k, Arrays.toString(getLeastNumbers(array, k))));
    }
}
